import torch
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset

N = 1000
batch_size = 32
seq_len = 5
input_size = 5
hidden_size = 8

rnn = nn.RNN(input_size=input_size, hidden_size=hidden_size, batch_first=True)
X = torch.randn(N, seq_len, input_size)

loader = DataLoader(TensorDataset(X), batch_size=batch_size, shuffle=True)

for i, (xb,) in enumerate(loader):
    out, h_n = rnn(xb)
    print(f"[RNN] batch {i}: xb={xb.shape} -> out={out.shape}, h_n={h_n.shape}")
    break
